from langchain_ollama import ChatOllama
from langchain.prompts import ChatPromptTemplate
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

# 初始化LLM
llm = ChatOllama(
    model="deepseek-r1:8b",
    base_url="http://localhost:11434",  # Ollama 服务地址
    temperature=0.7,  # 创造性程度
    num_predict=1024,  # 最大生成长度
    options={'stop': ['</think>', '</think>']} # 遇到这些标记就停止
)

# 创建一个简单的会话历史类
class SimpleChatMessageHistory(BaseChatMessageHistory):
    def __init__(self):
        self.messages = []
    
    def add_user_message(self, message):
        self.messages.append(HumanMessage(content=message))
    
    def add_ai_message(self, message):
        self.messages.append(AIMessage(content=message))
    
    def add_message(self, message):
        """添加单条消息"""
        self.messages.append(message)
    
    def add_messages(self, messages):
        """添加多条消息"""
        for message in messages:
            self.add_message(message)
    
    def clear(self):
        self.messages = []

# 创建聊天提示模板
chat_template = """The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.

{chat_history}
Human: {input}
AI:"""

chat_prompt = ChatPromptTemplate.from_template(chat_template)

# 创建基础chain
chain = chat_prompt | llm

# 模拟会话历史存储
store = {}

def get_session_history(session_id: str):
    """获取会话历史"""
    if session_id not in store:
        store[session_id] = SimpleChatMessageHistory()
    return store[session_id]

# 使用RunnableWithMessageHistory替代ConversationChain
conversation = RunnableWithMessageHistory(
    runnable=chain,
    get_session_history=get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
)

# 配置会话ID
config = {"configurable": {"session_id": "test-session"}}

# 第一天的对话
# 回合1
result = conversation.invoke({"input": "我姐姐明天要过生日，我需要一束生日花束。"}, config=config)
print("AI:", result.content if hasattr(result, 'content') else result)

# 打印第一次对话后的记忆
print("第一次对话后的记忆:")
for message in store["test-session"].messages:
    if isinstance(message, HumanMessage):
        print(f"Human: {message.content}")
    elif isinstance(message, AIMessage):
        print(f"AI: {message.content}")

        # 第一天的对话
# 回合2
result = conversation.invoke({"input": "她喜欢紫色和粉色。"}, config=config)
print("AI:", result.content if hasattr(result, 'content') else result)

# 打印第二次对话后的记忆
print("第二次对话后的记忆:")
for message in store["test-session"].messages:
    if isinstance(message, HumanMessage):
        print(f"Human: {message.content}")
    elif isinstance(message, AIMessage):
        print(f"AI: {message.content}")